# coding=utf-8
# 优化器 https://www.mindspore.cn/tutorials/zh-CN/master/advanced/modules/optimizer.html
# 若要为不同网络参数单独设置选项，如对卷积参数和非卷积参数设置不同的学习率，则可使用参数分组的方法来设置优化器。
# MindSpore中的trainable_params方法会屏蔽掉Parameter中requires_grad为False的属性，在为优化器配置 params 入参时，可使用net.trainable_params()方法来指定需要优化和更新的网络参数。
import numpy as np
import mindspore
from mindspore import nn, ops
from mindspore import Tensor, Parameter


class Net(nn.Cell):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(1, 6, 5, pad_mode="valid")
        self.param = Parameter(Tensor(np.array([1.0], np.float32)), 'param')

    def construct(self, x):
        x = self.conv(x)
        x = x * self.param
        out = ops.matmul(x, x)
        return out


net = Net()

# 配置优化器需要更新的参数
optim = nn.Adam(params=net.trainable_params())
print(net.trainable_params())

conv_params = [param for param in net.get_parameters() if 'conv' in param.name]
for conv_param in conv_params:
    conv_param.requires_grad = False
print(net.trainable_params())
optim = nn.Adam(params=net.trainable_params())

# 固定学习率：
# 设置学习率为0.01
optim = nn.Momentum(params=net.trainable_params(), learning_rate=0.01, momentum=0.9)
# 动态学习率：Dynamic LR
# 函数目前有基于余弦衰减函数计算学习率（nn.cosine_decay_lr）、基于指数衰减函数计算学习率（nn.exponential_decay_lr）、基于逆时衰减函数计算学习率（nn.inverse_decay_lr）、
# 基于自然指数衰减函数计算学习率（nn.natural_exp_decay_lr）、获取分段常量学习率（nn.piecewise_constant_lr）、基于多项式衰减函数计算学习率（nn.polynomial_decay_lr）和预热学习率(nn.warmup_lr)。
milestone = [1, 3, 10]
learning_rates = [0.1, 0.05, 0.01]
lr = nn.piecewise_constant_lr(milestone, learning_rates)

# 打印学习率
print(lr)

net = Net()
# 优化器设置待优化的网络参数和分段常量学习率
optim = nn.SGD(net.trainable_params(), learning_rate=lr)

# LearningRateSchedule类
# LearningRateSchedule类目前有基于余弦衰减函数计算学习率（nn.CosineDecayLR）、基于指数衰减函数计算学习率（nn.ExponentialDecayLR）、
# 基于逆时衰减函数计算学习率（nn.InverseDecayLR）、基于自然指数衰减函数计算学习率（nn.NaturalExpDecayLR）、
# 基于多项式衰减函数计算学习率（nn.PolynomialDecayLR）和预热学习率(nn.WarmUpLR)。
learning_rate = 0.1  # 学习率的初始值
decay_rate = 0.9  # 衰减率
decay_steps = 4  # 衰减的step数
step_per_epoch = 2

exponential_decay_lr = nn.ExponentialDecayLR(learning_rate, decay_rate, decay_steps)

for i in range(decay_steps):
    step = Tensor(i, mindspore.int32)
    result = exponential_decay_lr(step)
    print(f"step{i + 1}, lr:{result}")

net = Net()

# 优化器设置学习率为基于指数衰减函数计算学习率
optim = nn.Momentum(net.trainable_params(), learning_rate=exponential_decay_lr, momentum=0.9)

# 权重衰减:权重衰减(weight decay)，通常也被称为L2正则化，是一种缓解深度神经网络过拟合的方法。
net = Net()
optimizer = nn.Momentum(net.trainable_params(), learning_rate=0.01,
                        momentum=0.9, weight_decay=0.9)
# MindSpore还支持动态weight decay。此时weight_decay是用户自定义的一个Cell，称之为weight_decay_schedule。在训练过程中，优化器内部会调用该Cell的实例，传入global_step计算当前step的weight_decay值。
# 其中global_step是内部维护的变量，每训练一个step，global_step都会自加1。注意，自定义的weight_decay_schedule的construct仅接收一个输入。
from mindspore.nn import Cell
from mindspore import ops, nn
import mindspore as ms


class ExponentialWeightDecay(Cell):

    def __init__(self, weight_decay, decay_rate, decay_steps):
        super(ExponentialWeightDecay, self).__init__()
        self.weight_decay = weight_decay
        self.decay_rate = decay_rate
        self.decay_steps = decay_steps

    def construct(self, global_step):
        # construct只能有一个输入，训练过程中，会自动传入global step进行计算
        p = global_step / self.decay_steps
        return self.weight_decay * ops.pow(self.decay_rate, p)


net = Net()

weight_decay = ExponentialWeightDecay(weight_decay=0.0001, decay_rate=0.1, decay_steps=10000)
optimizer = nn.Momentum(net.trainable_params(), learning_rate=0.01,
                        momentum=0.9, weight_decay=weight_decay)

# 超参分组
# 下例分别对卷积参数和非卷积参数设置不同的学习率和权重衰减参数。
net = Net()

# 卷积参数
conv_params = list(filter(lambda x: 'conv' in x.name, net.trainable_params()))
# 非卷积参数
no_conv_params = list(filter(lambda x: 'conv' not in x.name, net.trainable_params()))

# 固定学习率
fix_lr = 0.01

# 基于多项式衰减函数计算学习率
polynomial_decay_lr = nn.PolynomialDecayLR(learning_rate=0.1,  # 学习率初始值
                                           end_learning_rate=0.01,  # 学习率最终值
                                           decay_steps=4,  # 衰减的step数
                                           power=0.5)  # 多项式幂

# 卷积参数使用固定学习率0.001,权重衰减为0.01
# 非卷积参数使用动态学习率，权重衰减为0.0
group_params = [{'params': conv_params, 'weight_decay': 0.01, 'lr': fix_lr},
                {'params': no_conv_params, 'lr': polynomial_decay_lr}]
# 当前MindSpore除个别优化器外（例如AdaFactor，FTRL），均支持对学习率进行分组，详情参考优化器API。
optim = nn.Momentum(group_params, learning_rate=0.1, momentum=0.9, weight_decay=0.0)


# 自定义优化器
# 自定义优化器时需继承优化器基类nn.Optimizer，并重写__init__方法和construct方法以自行设定参数更新策略。
class Momentum(nn.Optimizer):
    """定义优化器"""

    def __init__(self, params, learning_rate, momentum=0.9):
        super(Momentum, self).__init__(learning_rate, params)
        self.momentum = Parameter(Tensor(momentum, ms.float32), name="momentum")
        self.moments = self.parameters.clone(prefix="moments", init="zeros")

    def construct(self, gradients):
        """construct输入为梯度，在训练中自动传入梯度gradients"""
        lr = self.get_lr()
        params = self.parameters  # 待更新的权重参数

        for i in range(len(params)):
            # 更新moments值
            ops.assign(self.moments[i], self.moments[i] * self.momentum + gradients[i])
            update = params[i] - self.moments[i] * lr  # 带有动量的SGD算法
            ops.assign(params[i], update)
        return params


net = Net()
# 设置优化器待优化的参数和学习率为0.01
opt = Momentum(net.trainable_params(), 0.01)


# mindspore.ops也封装了优化器算子供用户自行定义优化器，
# 如ops.ApplyCenteredRMSProp、 ops.ApplyMomentum和ops.ApplyRMSProp等。下例使用ApplyMomentum算子自定义优化器Momentum：
class Momentum(nn.Optimizer):
    """定义优化器"""

    def __init__(self, params, learning_rate, momentum=0.9):
        super(Momentum, self).__init__(learning_rate, params)
        self.moments = self.parameters.clone(prefix="moments", init="zeros")
        self.momentum = momentum
        self.opt = ops.ApplyMomentum()

    def construct(self, gradients):
        # 待更新的权重参数
        params = self.parameters
        success = None
        for param, mom, grad in zip(params, self.moments, gradients):
            success = self.opt(param, mom, self.learning_rate, grad, self.momentum)
        return success


net = Net()
# 设置优化器待优化的参数和学习率为0.01
opt = Momentum(net.trainable_params(), 0.01)


# experimental.optim:
# 实验性的优化器模块 mindspore.experimental.optim，旨在对优化器做功能扩展。
# 除上述异同外，mindspore.experimental.optim 下的优化器还支持查看参数组、运行中修改优化器参数等功能，详见下文。


# 配置优化器
# 参数配置
# 常规场景下，与 mindspore.nn.optim 的参数配置方式相同，传入 net.trainable_params 即可。
# 学习率
# 固定学习率：
#
# 与 mindspore.nn.optim 的固定学习率配置方式相同。
#
# 动态学习率：
#
# mindspore.experimental.optim.lr_scheduler 下提供了动态学习率模块与 mindspore.experimental.optim 配合使用，使用方式与 mindspore.nn.optim 不同：
#
# mindspore.nn.optim：将动态学习率列表或实例传给优化器的入参 learning_rate，使用方式请参考DynamicLR函数和LearningRateSchedule类。
#
# mindspore.experimental.optim：将优化器实例传给动态学习率类的入参 optimizer，使用方式请参考LRScheduler类。
#
# LRScheduler 提供的获取学习率的方式：
#
# get_lr方法。以 StepLR 为例，训练过程中可以直接使用scheduler.get_lr() 手动获取学习率。

from mindspore.experimental import optim
net = Net()
optimizer = optim.Adam(net.trainable_params(), lr=0.1)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)
print(scheduler.get_last_lr())

# 超参分组
conv_params = list(filter(lambda x: 'conv' in x.name, net.trainable_params()))
no_conv_params = list(
    filter(lambda x: 'conv' not in x.name, net.trainable_params()))
group_params = [
    {'params': conv_params, 'weight_decay': 0.01, 'lr': 0.9, "amsgrad": True},
    {'params': no_conv_params, 'lr': 0.66, "eps": 1e-6, "betas": (0.8, 0.88)}]
optimizer = optim.Adam(params=group_params, lr=0.01)

# 查看优化器配置
print(optimizer.param_groups)
print(optimizer.param_groups[1]["lr"].value())
print(optimizer)

# 运行中修改优化器参数
net = Net()
loss_fn = nn.MAELoss()
optimizer = optim.Adam(net.trainable_params(), lr=0.1)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)
last_step_loss = 0.1

def forward_fn(data, label):
    logits = net(data)
    loss = loss_fn(logits, label)
    return loss

grad_fn = mindspore.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=True)

def train_step(data, label):
    (loss, _), grads = grad_fn(data, label)
    optimizer(grads)
    if ops.abs(loss - last_step_loss) < 0.1:
        ops.assign(optimizer.param_groups[1]["lr"], Tensor(0.01))
    return loss

# 运行中修改除lr以外的优化器参数: 目前仅PyNative模式下支持运行中修改其他优化器参数，Graph模式下的修改将不生效或报错。
net = Net()
loss_fn = nn.MAELoss()
optimizer = optim.Adam(net.trainable_params(), lr=0.1)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)
last_step_loss = 0.1

def forward_fn(data, label):
    logits = net(data)
    loss = loss_fn(logits, label)
    return loss

grad_fn = mindspore.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=True)

def train_step(data, label):
    (loss, _), grads = grad_fn(data, label)
    optimizer(grads)
    if ops.abs(loss - last_step_loss) < 0.1:
        optimizer.param_groups[1]["weight_decay"] = 0.02
    return loss
